import tushare as ts
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from mpl_finance import candlestick2_ohlc
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense

# 获取股票数据
pro = ts.pro_api('eb40c02798b5f92812c3bcc2cc837933c66ba60d3f63dfc1772b45ce')
df = pro.daily(ts_code='601939.SH', start_date='20100101', end_date='20241225')
df.to_csv("stock601939.csv")

# 数据预处理
mydata = pd.read_csv("stock601939.csv", parse_dates=["trade_date"], index_col="trade_date")[["open", "high", "low", "close"]]
mydata = mydata.sort_index(ascending=True)

# 绘制K线图和均线
fig, ax = plt.subplots(figsize=(14, 7))
candlestick2_ohlc(ax, opens=mydata['open'].values, highs=mydata['high'].values, lows=mydata['low'].values, closes=mydata['close'].values, width=0.5, colorup="r", colordown="g")

mydata['5'] = mydata['close'].rolling(window=5).mean()
mydata['10'] = mydata['close'].rolling(window=10).mean()

plt.plot(mydata['5'].values, alpha=0.5, label='MA5')
plt.plot(mydata['10'].values, alpha=0.5, label='MA10')
ax.legend(facecolor='white', edgecolor='white', fontsize=10)
plt.title('K线图及均线')
plt.xlabel('日期')
plt.ylabel('价格')
plt.grid(True)
plt.show()

# 划分数据集
training_set = mydata.iloc[:int(mydata.shape[0] * 0.8), 3:4].values  # 使用close价格
test_set = mydata.iloc[int(mydata.shape[0] * 0.8):, 3:4].values

# 数据归一化
sc = MinMaxScaler(feature_range=(0, 1))
training_set_scaled = sc.fit_transform(training_set)
testing_set_scaled = sc.transform(test_set)

# 准备数据集用于训练LSTM
n_timestamp = 60  # 过去60天的数据预测1天
def data_split(sequence, n_timestamp):
    X = []
    y = []
    for i in range(len(sequence) - n_timestamp):
        seq_x, seq_y = sequence[i:i+n_timestamp], sequence[i+n_timestamp]
        X.append(seq_x)
        y.append(seq_y)
    return np.array(X), np.array(y)

X_train, y_train = data_split(training_set_scaled, n_timestamp)
X_train = X_train.reshape(X_train.shape[0], X_train.shape[1], 1)
X_test, y_test = data_split(testing_set_scaled, n_timestamp)
X_test = X_test.reshape(X_test.shape[0], X_test.shape[1], 1)

# 构建LSTM模型
model = Sequential()
model.add(LSTM(units=50, activation='relu', input_shape=(X_train.shape[1], 1)))
model.add(Dense(units=1))

model.compile(optimizer='adam', loss='mean_squared_error')

# 训练模型
n_epochs = 50  # 设置epochs数量
history = model.fit(X_train, y_train, batch_size=64, epochs=n_epochs, validation_data=(X_test, y_test), validation_freq=1)

# 绘制损失曲线
plt.figure(figsize=(14, 7))
plt.plot(history.history['loss'], label='Training Loss')
plt.plot(history.history['val_loss'], label='Validation Loss')
plt.title('训练和验证损失')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend()
plt.grid(True)
plt.show()

# 预测股票价格
predicted_stock_price = model.predict(X_test)
predicted_stock_price = sc.inverse_transform(predicted_stock_price)
real_stock_price = sc.inverse_transform(y_test)

# 绘制预测结果
plt.figure(figsize=(14, 7))
plt.plot(real_stock_price, color='red', label='实际股票价格')
plt.plot(predicted_stock_price, color='blue', label='预测股票价格')
plt.title('股票价格预测')
plt.xlabel('时间')
plt.ylabel('股票价格')
plt.legend()
plt.grid(True)
plt.show()

# 计算评估指标
MSE = mean_squared_error(real_stock_price, predicted_stock_price)
RMSE = mean_squared_error(real_stock_price, predicted_stock_price) ** 0.5
MAE = mean_absolute_error(real_stock_price, predicted_stock_price)
R2 = r2_score(real_stock_price, predicted_stock_price)

print('均方误差: %.5f' % MSE)
print('均方根误差: %.5f' % RMSE)
print('平均绝对误差: %.5f' % MAE)
print('R2: %.5f' % R2)
import mpl_finance
